from dataclasses import dataclass
from typing import Literal
import pandas as pd
import numpy as np
from . import base


@dataclass
class ShadowLength(base.BarParams, base.BaseTransformer, base.CopyMixin):
    new_col: str = 'shadow_length'
    typ: Literal['upper', 'lower'] = 'upper'

    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        if self.typ == 'upper':
            a = data[self.col_high]
            b = data[[self.col_open, self.col_close]].max(axis=1)
        elif self.typ == 'lower':
            a = data[self.col_low]
            b = data[[self.col_open, self.col_close]].min(axis=1)
        else:
            raise ValueError('wrong `typ`')
        data[self.new_col] = (a - b).abs()  # type: ignore
        return data


@dataclass
class ShadowRatio(base.BarParams, base.BaseTransformer, base.CopyMixin):
    new_col: str = 'shadow_ratio'
    col_shadow_len: str = 'shadow_length'
    typ: Literal['upper', 'lower'] = 'upper'

    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        if self.col_shadow_len not in data:
            ncol = f'_shadow_{self.typ}'
            shadow_len = (ShadowLength(new_col=ncol,
                                       typ=self.typ,
                                       col_open=self.col_open,
                                       col_high=self.col_high,
                                       col_low=self.col_low,
                                       col_close=self.col_close)
                          .fit_transform(data)[ncol])
        else:
            shadow_len = data[self.col_shadow_len]
        tot_len = data[self.col_high] - data[self.col_low]
        data[self.new_col] = np.where((shadow_len > 0) & (tot_len > 0),
                                      shadow_len / tot_len, 0)  # 一字板的时候认为比例是0
        return data


@dataclass
class BodyLength(base.BarParams, base.BaseTransformer, base.CopyMixin):
    new_col: str = 'body_length'

    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        data[self.new_col] = (data[self.col_close] - data[self.col_open]).abs()
        return data


@dataclass
class BodyRatio(base.BarParams, base.BaseTransformer, base.CopyMixin):
    new_col: str = 'body_ratio'
    col_body_len: str = 'body_len'

    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        if self.col_body_len not in data:
            ncol = f'_body_len'
            body_len = (BodyLength(new_col=ncol,
                                   col_open=self.col_open,
                                   col_close=self.col_close)
                        .fit_transform(data)[ncol])
        else:
            body_len = data[self.col_body_len]
        tot_len = (data[self.col_high] - data[self.col_low])
        data[self.new_col] = np.where((body_len > 0) & (tot_len > 0), body_len / tot_len,
                                      np.where((body_len == 0) & (tot_len == 0), 1, 0))
        return data


@dataclass
class IsGap(base.BarParams, base.BaseTransformer, base.CopyMixin):
    new_col: str = 'is_gap'

    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        is_gap_up = (data[self.col_open] > data[self.col_close].shift(1)) & (data[self.col_close] > data[self.col_open])
        is_gap_down = (data[self.col_open] < data[self.col_close].shift(1)) & (
                data[self.col_close] < data[self.col_open])
        data[self.new_col] = np.where(is_gap_up | is_gap_down, 1, 0)
        return data


@dataclass
class GapMid(base.BarParams, base.BaseTransformer, base.CopyMixin):
    new_col: str = 'gap_mid'
    col_is_gap: str = 'is_gap'

    def transform(self, X: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
        data = self.copy_or(X)
        if self.col_is_gap not in data:
            ncol = f'_is_gap'
            is_gap = (IsGap(new_col=ncol,
                            col_open=self.col_open,
                            col_high=self.col_high,
                            col_low=self.col_low,
                            col_close=self.col_close)
                      .fit_transform(data))[ncol]
        else:
            is_gap = data[self.col_is_gap]
        idx = np.arange(len(data))
        idx_for = idx[is_gap == 1]
        idx_back = idx_for - 1
        idx_back[idx_back < 0] = 0
        data[self.new_col] = (data[self.col_close].iloc[idx_for] + data[self.col_close].iloc[
            idx_back].values) / 2  # type: ignore
        return data
